// Copyright (c) 2025 Shenzhen Kaihong Digital Industry Development Co., Ltd.
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::{collections::BTreeMap, path::Path};

use anyhow::Context;
use askama::Template;
use convert_case::{Case, Casing};

use crate::msg_def::{BaseType, Field, FieldType, MsgDef, MsgDefMap};

#[derive(Template)]
#[template(path = "lib.rs.askama")]
struct LibRsTemplate<'a> {
    map: &'a MsgDefMap,
}

#[derive(Template)]
#[template(path = "pkg_mod.rs.askama")]
struct PkgModTemplate<'a> {
    msg_defs: &'a BTreeMap<String, MsgDef>,
}

#[derive(Template)]
#[template(path = "msg.rs.askama", escape = "none")]
struct MsgRsTemplate<'a> {
    prefix: &'a str,
    msg_name: &'a str,
    fields: &'a [Field],
}

impl MsgDefMap {
    pub async fn gen_rust_crate<P: AsRef<Path>>(&self, out_dir: P) -> anyhow::Result<()> {
        let out_dir = out_dir.as_ref();
        let src_dir = out_dir.join("src");
        self.gen_rust_msg_mods(&src_dir, "crate").await
            .context("failed to generate Rust message modules")?;

        let lib_rs_path = src_dir.join("lib.rs");
        tokio::fs::write(&lib_rs_path, self.gen_lib_rs()?).await
            .context(format!("failed to write Rust file '{}'", lib_rs_path.display()))?;

        let cargo_toml_path = out_dir.join("Cargo.toml");
        tokio::fs::write(&cargo_toml_path, Self::gen_cargo_toml()).await
            .context(format!("failed to write Cargo.toml file '{}'", cargo_toml_path.display()))?;
        
        Ok(())
    }

    pub async fn gen_rust_mod<P: AsRef<Path>>(&self, out_dir: P, prefix: &str) -> anyhow::Result<()> {
        let out_dir = out_dir.as_ref();
        self.gen_rust_msg_mods(&out_dir, prefix).await
            .context("failed to generate Rust message modules")?;

        let mod_rs_path = out_dir.join("mod.rs");
        tokio::fs::write(&mod_rs_path, self.gen_lib_rs()?).await
            .context(format!("failed to write Rust file '{}'", mod_rs_path.display()))?;

        Ok(())
    }

    async fn gen_rust_msg_mods<P: AsRef<Path>>(&self, out_dir: P, prefix: &str) -> anyhow::Result<()> {
        let out_dir = out_dir.as_ref();
        for (pkg_name, msg_defs) in self.map.iter() {
            let pkg_dir = out_dir.join(pkg_name);
            tokio::fs::create_dir_all(&pkg_dir).await
                .context(format!("failed to create '{}'", pkg_dir.display()))?;

            for (msg_name, msg_def) in msg_defs.iter() {
                let content = msg_def.gen_msg_rs(prefix, msg_name)
                    .context(format!("failed to generate Rust file for message {}/{}", pkg_name, msg_name))?;

                let path = pkg_dir.join(format!("{}.rs", msg_name.to_case(Case::Snake)));
                tokio::fs::write(&path, content).await
                    .context(format!("failed to write Rust file '{}'", path.display()))?;
            }

            let mod_path = pkg_dir.join("mod.rs");
            tokio::fs::write(&mod_path, Self::gen_pkg_mod(msg_defs)?).await
                .context(format!("failed to write Rust file '{}'", mod_path.display()))?;
        }
        Ok(())
    }

    fn gen_cargo_toml() -> String {
        r#"[package]
name = "_msgs"
version = "0.1.0"
edition = "2024"

[dependencies]
ohmw-core = { path = "../../../../framework/core" }
ohmw-derive = { path = "../../../../framework/derive" }
ohmw-message = { path = "../../../../framework/message" }
serde = { version = "1", features = ["derive"] }
"#
       .to_string()
    }

    fn gen_lib_rs(&self) -> anyhow::Result<String> {
        Ok(LibRsTemplate { map: self }.render()?)
    }

    fn gen_pkg_mod(msg_defs: &BTreeMap<String, MsgDef>) -> anyhow::Result<String> {
        Ok(PkgModTemplate { msg_defs }.render()?)
    }
}

impl MsgDef {
    fn gen_msg_rs(&self, prefix: &str, msg_name: &str) -> anyhow::Result<String> {
        Ok(MsgRsTemplate {
            prefix,
            msg_name,
            fields: &self.fields,
        }.render()?)
    }
}

impl Field {
    fn field_name_to_rust(&self) -> String {
        let s = self.field_name.to_case(Case::Snake);
        if syn::parse_str::<syn::Ident>(s.as_str()).is_err() {
            format!("r#{}", s)
        } else {
            s
        }
    }

    fn field_type_to_rust(&self, prefix: &str, msg_name: &str) -> String {
        let mut field_type = self.field_type.to_rust_type(prefix);
        if msg_name.to_case(Case::Pascal) == field_type && field_type == "String" {
            field_type = String::from("std::string::String");
        }
        field_type
    }
}

impl FieldType {
    fn to_rust_type(&self, prefix: &str) -> String {
        if self.is_array {
            if let Some(size) = self.array_size {
                format!("[{}; {}]", self.base_type.to_rust_type(prefix), size)
            } else {
                format!("Vec<{}>", self.base_type.to_rust_type(prefix))
            }
        } else {
            self.base_type.to_rust_type(prefix)
        }
    }
}

impl BaseType {
    fn to_rust_type(&self, prefix: &str) -> String {
        match (self.pkg_name(), self.msg_name()) {
            (Some(pkg_name), msg_name) => format!("{}::{}::{}",
                prefix,
                pkg_name,
                msg_name.to_case(Case::Pascal),
            ),
            (None, msg_name) => match msg_name {
                "bool" => "bool".to_string(),
                "byte" => "u8".to_string(),
                "char" => "u8".to_string(),
                "int8" => "i8".to_string(),
                "uint8" => "u8".to_string(),
                "int16" => "i16".to_string(),
                "uint16" => "u16".to_string(),
                "int32" => "i32".to_string(),
                "uint32" => "u32".to_string(),
                "int64" => "i64".to_string(),
                "uint64" => "u64".to_string(),
                "float32" => "f32".to_string(),
                "float64" => "f64".to_string(),
                "string" => "String".to_string(),
                "time" => "ohmw_message::Time".to_string(),
                "duration" => "ohmw_message::Duration".to_string(),
                other => format!("not_found::{}", other.to_case(Case::Pascal)),
            },
        }
    }
}
